{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53",
   "metadata": {},
   "source": [
    "# How to add thread-level persistence (functional API)\n",
    "\n",
    "!!! info \"Prerequisites\"\n",
    "\n",
    "    This guide assumes familiarity with the following:\n",
    "    \n",
    "    - [Functional API](../../concepts/functional_api/)\n",
    "    - [Persistence](../../concepts/persistence/)\n",
    "    - [Memory](../../concepts/memory/)\n",
    "    - [Chat Models](https://python.langchain.com/docs/concepts/chat_models/)\n",
    "\n",
    "!!! info \"Not needed for LangGraph API users\"\n",
    "\n",
    "    If you're using the LangGraph API, you needn't manually implement a checkpointer. The API automatically handles checkpointing for you. This guide is relevant when implementing LangGraph in your own custom server.\n",
    "\n",
    "Many AI applications need memory to share context across multiple interactions on the same [thread](../../concepts/persistence#threads) (e.g., multiple turns of a conversation). In LangGraph functional API, this kind of memory can be added to any [entrypoint()][langgraph.func.entrypoint] workflow using [thread-level persistence](https://langchain-ai.github.io/langgraph/concepts/persistence).\n",
    "\n",
    "When creating a LangGraph workflow, you can set it up to persist its results by using a [checkpointer](https://langchain-ai.github.io/langgraph/reference/checkpoints/#basecheckpointsaver):\n",
    "\n",
    "\n",
    "1. Create an instance of a checkpointer:\n",
    "\n",
    "    ```python\n",
    "    from langgraph.checkpoint.memory import InMemorySaver\n",
    "    \n",
    "    checkpointer = InMemorySaver()       \n",
    "    ```\n",
    "\n",
    "2. Pass `checkpointer` instance to the `entrypoint()` decorator:\n",
    "\n",
    "    ```python\n",
    "    from langgraph.func import entrypoint\n",
    "    \n",
    "    @entrypoint(checkpointer=checkpointer)\n",
    "    def workflow(inputs)\n",
    "        ...\n",
    "    ```\n",
    "\n",
    "3. Optionally expose `previous` parameter in the workflow function signature:\n",
    "\n",
    "    ```python\n",
    "    @entrypoint(checkpointer=checkpointer)\n",
    "    def workflow(\n",
    "        inputs,\n",
    "        *,\n",
    "        # you can optionally specify `previous` in the workflow function signature\n",
    "        # to access the return value from the workflow as of the last execution\n",
    "        previous\n",
    "    ):\n",
    "        previous = previous or []\n",
    "        combined_inputs = previous + inputs\n",
    "        result = do_something(combined_inputs)\n",
    "        ...\n",
    "    ```\n",
    "\n",
    "4. Optionally choose which values will be returned from the workflow and which will be saved by the checkpointer as `previous`:\n",
    "\n",
    "    ```python\n",
    "    @entrypoint(checkpointer=checkpointer)\n",
    "    def workflow(inputs, *, previous):\n",
    "        ...\n",
    "        result = do_something(...)\n",
    "        return entrypoint.final(value=result, save=combine(inputs, result))\n",
    "    ```\n",
    "\n",
    "This guide shows how you can add thread-level persistence to your workflow.\n",
    "\n",
    "!!! tip \"Note\"\n",
    "\n",
    "    If you need memory that is __shared__ across multiple conversations or users (cross-thread persistence), check out this [how-to guide](../cross-thread-persistence-functional).\n",
    "\n",
    "!!! tip \"Note\"\n",
    "\n",
    "    If you need to add thread-level persistence to a `StateGraph`, check out this [how-to guide](../persistence)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cbd446a-808f-4394-be92-d45ab818953c",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First we need to install the packages required"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "af4ce0ba-7596-4e5f-8bf8-0b0bd6e62833",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install --quiet -U langgraph langchain_anthropic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0abe11f4-62ed-4dc4-8875-3db21e260d1d",
   "metadata": {},
   "source": [
    "Next, we need to set API key for Anthropic (the LLM we will use)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c903a1cf-2977-4e2d-ad7d-8b3946821d89",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "\n",
    "def _set_env(var: str):\n",
    "    if not os.environ.get(var):\n",
    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
    "\n",
    "\n",
    "_set_env(\"ANTHROPIC_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0ed46a8-effe-4596-b0e1-a6a29ee16f5c",
   "metadata": {},
   "source": [
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Set up <a href=\"https://smith.langchain.com\">LangSmith</a> for LangGraph development</p>\n",
    "    <p style=\"padding-top: 5px;\">\n",
    "        Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started <a href=\"https://docs.smith.langchain.com\">here</a>. \n",
    "    </p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cf509bc",
   "metadata": {},
   "source": [
    "## Example: simple chatbot with short-term memory\n",
    "\n",
    "We will be using a workflow with a single task that calls a [chat model](https://python.langchain.com/docs/concepts/chat_models/).\n",
    "\n",
    "Let's first define the model we'll be using:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "892b54b9-75f0-4804-9ed0-88b5e5532989",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_anthropic import ChatAnthropic\n",
    "\n",
    "model = ChatAnthropic(model=\"claude-3-5-sonnet-latest\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b7a2792-982b-4e47-83eb-0c594725d1c1",
   "metadata": {},
   "source": [
    "Now we can define our task and workflow. To add in persistence, we need to pass in a [Checkpointer](https://langchain-ai.github.io/langgraph/reference/checkpoints/#langgraph.checkpoint.base.BaseCheckpointSaver) to the [entrypoint()][langgraph.func.entrypoint] decorator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "87326ea6-34c5-46da-a41f-dda26ef9bd74",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import BaseMessage\n",
    "from langgraph.graph import add_messages\n",
    "from langgraph.func import entrypoint, task\n",
    "from langgraph.checkpoint.memory import InMemorySaver\n",
    "\n",
    "\n",
    "@task\n",
    "def call_model(messages: list[BaseMessage]):\n",
    "    response = model.invoke(messages)\n",
    "    return response\n",
    "\n",
    "\n",
    "checkpointer = InMemorySaver()\n",
    "\n",
    "\n",
    "@entrypoint(checkpointer=checkpointer)\n",
    "def workflow(inputs: list[BaseMessage], *, previous: list[BaseMessage]):\n",
    "    if previous:\n",
    "        inputs = add_messages(previous, inputs)\n",
    "\n",
    "    response = call_model(inputs).result()\n",
    "    return entrypoint.final(value=response, save=add_messages(inputs, response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "250d8fd9-2e7a-4892-9adc-19762a1e3cce",
   "metadata": {},
   "source": [
    "If we try to use this workflow, the context of the conversation will be persisted across interactions:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7654ebcc-2179-41b4-92d1-6666f6f8634f",
   "metadata": {},
   "source": [
    "!!! note Note\n",
    "\n",
    "    If you're using LangGraph Platform or LangGraph Studio, you __don't need__ to pass checkpointer to the entrypoint decorator, since it's done automatically."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a1b56c5-bd61-4192-8bdb-458a1e9f0159",
   "metadata": {},
   "source": [
    "We can now interact with the agent and see that it remembers previous messages!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cfd140f0-a5a6-4697-8115-322242f197b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Hi Bob! I'm Claude. Nice to meet you! How are you today?\n"
     ]
    }
   ],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "input_message = {\"role\": \"user\", \"content\": \"hi! I'm bob\"}\n",
    "for chunk in workflow.stream([input_message], config, stream_mode=\"values\"):\n",
    "    chunk.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bb07bf8-68b7-4049-a0f1-eb67a4879a3a",
   "metadata": {},
   "source": [
    "You can always resume previous threads:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "08ae8246-11d5-40e1-8567-361e5bef8917",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Your name is Bob.\n"
     ]
    }
   ],
   "source": [
    "input_message = {\"role\": \"user\", \"content\": \"what's my name?\"}\n",
    "for chunk in workflow.stream([input_message], config, stream_mode=\"values\"):\n",
    "    chunk.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f47bbfc-d9ef-4288-ba4a-ebbc0136fa9d",
   "metadata": {},
   "source": [
    "If we want to start a new conversation, we can pass in a different `thread_id`. Poof! All the memories are gone!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "273d56a8-f40f-4a51-a27f-7c6bb2bda0ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "I don't know your name unless you tell me. Each conversation I have starts fresh, so I don't have access to any previous interactions or personal information unless you share it with me.\n"
     ]
    }
   ],
   "source": [
    "input_message = {\"role\": \"user\", \"content\": \"what's my name?\"}\n",
    "for chunk in workflow.stream(\n",
    "    [input_message],\n",
    "    {\"configurable\": {\"thread_id\": \"2\"}},\n",
    "    stream_mode=\"values\",\n",
    "):\n",
    "    chunk.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac7926a8-4c88-4b16-973c-53d6da3f4a08",
   "metadata": {},
   "source": [
    "!!! tip \"Streaming tokens\"\n",
    "\n",
    "    If you would like to stream LLM tokens from your chatbot, you can use `stream_mode=\"messages\"`. Check out this [how-to guide](../streaming-tokens) to learn more."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
